/*
 * Copyright 2022 The Open Islands Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.openislands.oi.scheduling;

import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.conditions.update.UpdateWrapper;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.core.toolkit.support.SFunction;
import com.google.common.collect.Sets;
import org.openislands.oi.config.stateengine.StateEngineContext;
import org.openislands.oi.constant.*;
import org.openislands.oi.dao.entity.NodeDO;
import org.openislands.oi.dao.entity.OfNodeTaskDO;
import org.openislands.oi.dao.entity.JobDO;
import org.openislands.oi.dao.entity.TaskDO;
import org.openislands.oi.error.LogicException;
import org.openislands.oi.pojo.dto.CommonResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

@Component
public class StateEngine implements InitializingBean {
    private static final Logger log = LoggerFactory.getLogger(StateEngine.class);
    // rules
    private static final Set<JobState> ALL_JOB_STATES = new HashSet<>(Arrays.asList(JobState.class.getEnumConstants()));
    private Set<JobState> INTERRUPT_JOB_STATES;
    private Set<JobState> FINISHED_JOB_STATES;
    private Set<JobState> UNFINISHED_JOB_STATES;

    private final static Set<TaskState> ALL_TASK_STATES = new HashSet<>(Arrays.asList(TaskState.class.getEnumConstants()));
    private Set<TaskState> INTERRUPT_TASK_STATES;

    private Set<TaskState> FINISHED_TASK_STATES;
    private Set<TaskState> UNFINISHED_TASK_STATES;

    private Set<NodeState> CAN_DELETE_NODE_STATES;
    private Set<NodeState> COOPERATION_NODE_STATES;

    private Map<TaskState, JobState> TASK_STATE_EXTRA_MAP;

    /***
     * <K do class>
     * <V corresponding state enum>
     */
    private final Map<Class<?>, StateRuleEnum> doRuleCache = new HashMap<>();

    @Resource
    private StateEngineContext stateEngineContext;

    @Override
    public void afterPropertiesSet() {
        initializeRules();
    }

    public void initializeRules() {
        // add the corresponding state of the entity class
        doRuleCache.put(JobDO.class, StateRuleEnum.JOB);
        doRuleCache.put(TaskDO.class, StateRuleEnum.TASK);
        doRuleCache.put(OfNodeTaskDO.class, StateRuleEnum.TASK);
        doRuleCache.put(NodeDO.class, StateRuleEnum.NODE);

        INTERRUPT_JOB_STATES = stateEngineContext.getStateSet(Dict.INTERRUPT_JOB_STATES, JobState.class, JobState::name);
        FINISHED_JOB_STATES = stateEngineContext.getStateSet(Dict.FINISHED_JOB_STATES, JobState.class, JobState::name);
        UNFINISHED_JOB_STATES = Sets.difference(ALL_JOB_STATES, FINISHED_JOB_STATES);

        INTERRUPT_TASK_STATES = stateEngineContext.getStateSet(Dict.INTERRUPT_TASK_STATES, TaskState.class, TaskState::name);
        FINISHED_TASK_STATES = stateEngineContext.getStateSet(Dict.FINISHED_TASK_STATES, TaskState.class, TaskState::name);
        UNFINISHED_TASK_STATES = Sets.difference(ALL_TASK_STATES, FINISHED_TASK_STATES);

        CAN_DELETE_NODE_STATES = stateEngineContext.getStateSet(Dict.CAN_DELETE_NODE_STATES, NodeState.class, NodeState::name);
        COOPERATION_NODE_STATES = stateEngineContext.getStateSet(Dict.COOPERATION_NODE_STATES, NodeState.class, NodeState::name);

        //TASK_STATE_EXTRA_MAP.put(TaskState.READY, JobState.SUBMITTED);

        log.info("state engine: \n{}", this.stateEngineContext.toString());
    }

    public Boolean isInterruptedJobState(JobState jobState) {
        return this.INTERRUPT_JOB_STATES.contains(jobState);
    }

    public Boolean isFinishedJobState(JobState jobState) {
        return this.FINISHED_JOB_STATES.contains(jobState);
    }

    public Boolean isInterruptedTaskState(TaskState taskState) {
        return INTERRUPT_TASK_STATES.contains(taskState);
    }

    public Boolean isFinishedTaskState(TaskState taskState) {
        return FINISHED_TASK_STATES.contains(taskState);
    }

    public Boolean isUnfinishedTaskState(TaskState taskState) {
        return UNFINISHED_TASK_STATES.contains(taskState);
    }

    public Boolean isCanDeleteNodeState(NodeState nodeState) {
        return CAN_DELETE_NODE_STATES.contains(nodeState);
    }

    public boolean isCooperationNodeState(NodeState state) {
        return COOPERATION_NODE_STATES.contains(state);
    }

    public static Set<JobState> getAllJobStates() {
        return ALL_JOB_STATES;
    }

    public Set<JobState> getInterruptJobStates() {
        return INTERRUPT_JOB_STATES;
    }

    public Set<JobState> getUnfinishedJobStates() {
        return UNFINISHED_JOB_STATES;
    }

    public List<TaskState> intersectionWithTaskStates(Set<TaskState> a, Set<TaskState> b) {
        Set<TaskState> r = Sets.intersection(a, b);
        if (r.isEmpty()) {
            return new ArrayList<>();
        } else {
            return Arrays.stream(r.toArray(new TaskState[0])).sorted(Comparator.comparingInt(TaskState::getCode).reversed()).collect(Collectors.toList());
        }
    }

    public Set<TaskState> getInterruptTaskStates() {
        return INTERRUPT_TASK_STATES;
    }

    public Set<TaskState> getUnfinishedTaskStates() {
        return UNFINISHED_TASK_STATES;
    }

    public JobState getTaskStateMapJob(TaskState taskState) {
        try {
            return JobState.valueOf(taskState.name());
        } catch (Exception e) {
            if (this.isInterruptedTaskState(taskState)) {
                return JobState.RUN_FAILED;
            }
            if (this.isUnfinishedTaskState(taskState)) {
                return JobState.SUBMITTED;
            }
            throw new LogicException(MessageCodeEnum.DAG_SCHEDULING_ERROR, String.format("can not map task state %s to job state", taskState));
        }
    }

    /***
     * update state based on state engine rules
     * @param baseMapper        DO mapper
     * @param toState           JobState or TaskState
     * @param updateField       need update field function. <code> DO::getXXX </code>
     * @param getStateString    function to get status string. <code> State.byCode(DO::getCode).toString() </code>
     * @param wrapperFunction   additional required by selectOne where conditions
     * @param <DO>              DO type
     * @param <Field>           DO entity field type
     * @return update rows
     */
    public <DO, Field> CommonResult<Object> updateState(BaseMapper<DO> baseMapper, Field toState,
                                                        SFunction<DO, Field> updateField,
                                                        Function<Field, String> getStateString,
                                                        Function<LambdaUpdateWrapper<DO>, LambdaUpdateWrapper<DO>> wrapperFunction) {
        try {
            // query one
            LambdaUpdateWrapper<DO> wrapper = wrapperFunction.apply(new UpdateWrapper<DO>().lambda());
            final DO aDo = baseMapper.selectOne(wrapper);
            LogicException.nonNull(aDo, MessageCodeEnum.DATA_NOT_FOUND_ERROR);

            Field oldState = updateField.apply(aDo);
            String oldStateName = getStateString.apply(oldState);
            String toStateName = getStateString.apply(toState);

            if (Objects.equals(oldStateName, toStateName)) {
                return CommonResult.error(MessageCodeEnum.STATE_UPDATE_NOT_EFFECTIVE);
            }

            // transition check
            StateRuleEnum stateRuleEnum = doRuleCache.get(aDo.getClass());
            boolean condition = stateEngineContext.transition(stateRuleEnum, oldStateName, toStateName);
            if (!condition) {
                return new CommonResult<>(MessageCodeEnum.STATE_TRANSITION_FAILED, String.format("%s current state %s can only transition to %s states",
                        stateRuleEnum, oldStateName, stateEngineContext.getTransitionList(stateRuleEnum, oldStateName)));
            }

            // update
            int rows = baseMapper.update(null, wrapper.set(updateField, toState).eq(updateField, oldState));
            if (rows > 0) {
                return CommonResult.success();
            } else {
                return CommonResult.error(MessageCodeEnum.STATE_UPDATE_NOT_EFFECTIVE);
            }
        } catch (LogicException logicException) {
            throw logicException;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new LogicException(e, MessageCodeEnum.STATE_UPDATE_ERROR, e.getMessage());
        }
    }
}
